import os
import time
import logging

import torch

from amber import arguments
from amber.utils import redir_output
from amber.simulator import Simulator

if __name__ == "__main__":
    device = "cuda:0" if torch.cuda.is_available() else "cpu"

    args = arguments.init()
    if args.verbose >= 2:
        logging.basicConfig(level=logging.DEBUG)
    elif args.verbose >= 1:
        logging.basicConfig(level=logging.INFO)
    else:
        logging.basicConfig(level=logging.ERROR)

    logfile = os.path.join(args.log_dir, f"main-{time.strftime('%Y%m%d-%H%M%S')}.log")
    redir_output(logfile)

    sim = Simulator(
        prot_name=args.prot_name,
        prot_path=args.prot_file,
        utils_dir=args.utils_dir,
        command_save_path=os.path.join(args.log_dir, "PreprocessBackup"),
        log_dir=args.log_dir,
        temp_k=args.temp_k,
        device=device,
    )

    sim.run_preprocess()
    for i in range(5):
        sim.run_simulate(path=f'simulation-{i}', random_seed=3407*i)
    sim.run_postprocess()

